from PyQt5 import uic
from PyQt5.QtWidgets import QApplication, QMessageBox
from PyQt5.QtGui import QStandardItemModel
from PyQt5.QtCore import Qt
import pyautogui
import os

UI_path = "ui.ui"

DEBUG_PRINT = False

INIT_ROW_NUM = 20
INIT_COL_NUM = 2

CONFIDENCE = 0.9

DOSE_FOLDER = "dose screenshot"
EFFECT_FOLDER = "effect screenshot"


def return_files(target_path, ext='.png'):
    results = []
    for filename in os.listdir(target_path):
        file_ext = os.path.splitext(filename)[1]
        if file_ext == ext:
            result = os.path.join(target_path, filename)
            results.append(result)
    return results


def click_input(x, y, content):
    if isinstance(content, int) or isinstance(content, float):
        content = str(content)
    pyautogui.click(x, y)
    pyautogui.typewrite(content)


class MainWindow:

    def __init__(self):
        # 加载UI
        self.ui = uic.loadUi(UI_path)
        model = QStandardItemModel(INIT_ROW_NUM, INIT_COL_NUM)
        self.ui.tableView.setModel(model)
        # 连接信号
        self.ui.run.clicked.connect(self._run)

    def _locate(self, image_path):
        try:
            location = pyautogui.locateOnScreen(
                image_path, grayscale=True, confidence=CONFIDENCE)
            if location is not None:
                center = pyautogui.center(location)
                return center
            else:
                return None, None
        except OSError:
            QMessageBox.critical(
                self.ui, "Error", "图片 \""+image_path+"\" 的路径中不能含有中文！")
            return None, None

    def _run(self):
        if DEBUG_PRINT:
            print("# ---------- start running ---------- #")
        # 改变鼠标样式
        self.ui.setCursor(Qt.WaitCursor)
        # 数据获取
        data = self._get_data()  # list,[[d0,e0],[d1,e1],...]
        if data is None:
            self.ui.setCursor(Qt.ArrowCursor)
            return
        # 查找输入框位置
        dose_paths = return_files(DOSE_FOLDER)
        effect_paths = return_files(EFFECT_FOLDER)
        dose_x, dose_y, effect_x, effect_y = None, None, None, None
        for dose_path in dose_paths:
            x, y = self._locate(dose_path)
            if DEBUG_PRINT:
                print(dose_path)
                print(x, y)
            if x is not None:
                dose_x, dose_y = x, y
                break
        for effect_path in effect_paths:
            x, y = self._locate(effect_path)
            if DEBUG_PRINT:
                print(effect_path)
                print(x, y)
            if x is not None:
                effect_x, effect_y = x, y
                break
        # 如果成功查找到位置则进行界面操作，否则弹出警告
        if dose_x is None and effect_x is None:
            QMessageBox.critical(self.ui, "Error", "未找到 dose 输入框与 effect 输入框！")
        elif dose_x is None:
            QMessageBox.critical(self.ui, "Error", "未找到 dose 输入框!")
        elif effect_x is None:
            QMessageBox.critical(self.ui, "Error", "未找到 effect 输入框！")
        else:
            for d in data:
                click_input(dose_x, dose_y, d[0])
                pyautogui.press("enter")
                click_input(effect_x, effect_y, d[1])
                pyautogui.press("enter")
        # 还原鼠标样式
        self.ui.setCursor(Qt.ArrowCursor)

    def _get_data(self):
        """获取表格内容"""
        datas = []
        model = self.ui.tableView.model()
        if model.data(model.index(0, 0)) is None or model.data(model.index(0, 0)) == '':
            QMessageBox.critical(self.ui, "Error", "表格的第一行不能为空！")
            return None
        flag = False
        for row in range(model.rowCount()):
            datas.append([])
            for col in range(model.columnCount()):
                index = model.index(row, col)
                data = model.data(index)
                if data is not None:
                    datas[row].append(data)
                else:  # 第一个空行就退出
                    flag = True
                    break
            if flag:
                datas.pop()
                break
        return datas

    def _check(self, location, warning="未找到指定窗口，请保证指定窗口在屏幕上层！"):
        if location is None:
            QMessageBox.critical(self.ui, "Error", warning)
            return False
        else:
            return True


if __name__ == '__main__':
    app = QApplication([])
    main_window = MainWindow()
    main_window.ui.show()
    app.exec_()
